import akshare as ak
import time
import random
from database_manager import get_db_connection
from stock_data_provider import process_stock_symbol

def fetch_all_stocks():
    """获取所有A股股票信息并存储到数据库"""
    print("开始获取所有A股股票信息...")
    
    # 获取连接
    conn = get_db_connection()
    if not conn:
        print("数据库连接失败，无法继续")
        return
    
    try:
        # 使用akshare获取A股股票列表
        print("从akshare获取A股股票列表...")
        stock_df = ak.stock_zh_a_spot()
        print(f"成功获取 {len(stock_df)} 只股票")
        
        # 准备批量插入数据
        stocks_to_insert = []
        for _, row in stock_df.iterrows():
            # 获取标准化的股票代码和带后缀的完整代码
            symbol = row['代码']
            processed_symbol = process_stock_symbol(symbol)
            stock_name = row['名称']
            
            # 行业信息暂时留空，可以后续补充
            industry = None
            
            stocks_to_insert.append((processed_symbol, symbol, stock_name, industry))
        
        # 批量插入到数据库
        print("开始将股票信息插入数据库...")
        with conn.cursor() as cursor:
            # 清空现有表数据
            cursor.execute("TRUNCATE TABLE stock_info")
            
            # 批量插入
            sql = """
                INSERT INTO stock_info 
                (symbol, symbol_standardized, stock_name, industry) 
                VALUES (%s, %s, %s, %s)
            """
            
            # 分批插入，避免数据量过大
            batch_size = 500
            total_inserted = 0
            
            for i in range(0, len(stocks_to_insert), batch_size):
                batch = stocks_to_insert[i:i+batch_size]
                cursor.executemany(sql, batch)
                conn.commit()
                total_inserted += len(batch)
                print(f"已插入 {total_inserted}/{len(stocks_to_insert)} 条股票数据")
                
                # 添加小延迟，避免请求过于频繁
                if i + batch_size < len(stocks_to_insert):
                    time.sleep(0.5)
        
        print(f"所有股票数据插入完成，共 {total_inserted} 条记录")
        
    except Exception as e:
        print(f"获取或存储股票信息失败: {e}")
        conn.rollback()
    finally:
        conn.close()
        
    # 尝试补充行业信息
    print("\n开始尝试补充行业信息...")
    supplement_industry_info()

def supplement_industry_info():
    """尝试补充股票的行业信息"""
    conn = get_db_connection()
    if not conn:
        return
    
    try:
        # 获取没有行业信息的股票
        with conn.cursor() as cursor:
            cursor.execute("SELECT id, symbol_standardized FROM stock_info WHERE industry IS NULL LIMIT 200")
            stocks_to_update = cursor.fetchall()
            
        if not stocks_to_update:
            print("所有股票都已有行业信息或已尝试补充")
            return
        
        print(f"开始补充 {len(stocks_to_update)} 只股票的行业信息...")
        
        # 使用akshare获取行业信息
        for i, stock in enumerate(stocks_to_update):
            stock_id = stock['id']
            symbol = stock['symbol_standardized']
            
            try:
                # 使用行业分类接口
                try:
                    # 方式1：使用股票行业分类接口
                    industry_df = ak.stock_board_industry_name_ths()
                    # 检查列名，根据akshare的最新版本调整
                    if 'code' in industry_df.columns and 'name' in industry_df.columns:
                        # 查找对应的行业
                        industry_row = industry_df[industry_df['code'] == symbol]
                        if not industry_row.empty:
                            industry = industry_row.iloc[0]['name']
                            
                            # 更新数据库
                            with conn.cursor() as cursor:
                                cursor.execute("UPDATE stock_info SET industry = %s WHERE id = %s", (industry, stock_id))
                            conn.commit()
                            print(f"已补充 {symbol} 的行业信息: {industry}")
                    elif '代码' in industry_df.columns and '行业' in industry_df.columns:
                        # 兼容旧版本的列名
                        industry_row = industry_df[industry_df['代码'] == symbol]
                        if not industry_row.empty:
                            industry = industry_row.iloc[0]['行业']
                            
                            # 更新数据库
                            with conn.cursor() as cursor:
                                cursor.execute("UPDATE stock_info SET industry = %s WHERE id = %s", (industry, stock_id))
                            conn.commit()
                            print(f"已补充 {symbol} 的行业信息: {industry}")
                except:
                    # 如果方式1失败，记录但继续
                    pass
                
                # 随机延迟避免请求过于频繁
                time.sleep(random.uniform(0.5, 1.5))
                
                # 每处理20只股票，休息一小会
                if (i + 1) % 20 == 0:
                    print(f"已处理 {i + 1} 只股票，休息2秒...")
                    time.sleep(2)
                    
            except Exception as e:
                print(f"获取 {symbol} 行业信息失败: {e}")
                # 出错时稍作延迟
                time.sleep(1)
                
    except Exception as e:
        print(f"补充行业信息过程中出错: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    fetch_all_stocks()
